import pymysql
from common.yaml_utils import yaml_utils


class db:
    def __init__(self):
        mysql_utils = yaml_utils().get_mysql_config()
        self.host = mysql_utils["host"]
        self.port = mysql_utils["port"]
        self.user = mysql_utils["user"]
        self.password = mysql_utils["password"]
        self.database = mysql_utils["database"]
        self.cursor = None
        self.conn = None

    def connect(self):
        """建立数据库连接，返回是否成功"""
        try:
            self.conn = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset="utf8"
            )
            self.cursor = self.conn.cursor()
            return True
        except Exception as e:
            print(f"数据库连接失败: {e}")
            return False

    def close(self):
        """安全关闭连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        self.cursor = None
        self.conn = None

    def query(self, sql):
        """执行查询并返回结果（列表），失败返回None"""
        if not self.connect():  # 检查连接是否成功
            return None

        try:
            self.cursor.execute(sql)
            result = self.cursor.fetchall()  # 获取查询结果（列表）
            return result if result else None
        except Exception as e:
            print(f"查询失败: {e}")
            return None
        finally:
            self.close()  # 确保无论如何都会关闭连接